import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.init as init

num_feat = 153
num_data = 1000
num_epoch = 1000
x = init.uniform_(torch.Tensor(num_data, num_feat), -10, 10)
y = torch.mean(x * 5 - x**2 * 0.01, 1, True)

model = nn.Sequential(
          nn.Linear(num_feat, 100),
          nn.ReLU(),
          nn.Linear(100,100),
          nn.ReLU(),
          nn.Linear(100,100),
          nn.ReLU(),
          nn.Linear(100,1),
      )
loss_func = nn.MSELoss() 
optimizer = optim.SGD(model.parameters(), lr=0.0002)

# train
loss_array = []
for i in range(num_epoch):
    optimizer.zero_grad()
    output = model(x)
    loss = loss_func(output, y)
    loss.backward()
    optimizer.step()
    loss_array.append(loss)


import matplotlib.pyplot as plt
plt.plot(loss_array)
plt.show()

plt.figure(figsize=(10,10))
# plt.scatter(x.detach().numpy(),y,label="Original Data")
# plt.scatter(x.detach().numpy(),output.detach().numpy(),label="Model Output")
plt.scatter(y, output.detach().numpy(),label="Model Output")
plt.legend()
plt.show()

# save
test_sample_input_trace = torch.tensor(x)
test_sample_trace = torch.jit.trace(model, x)

test_sample_trace.save("example.pt")
